package org.dicr.sys.linux;

import java.io.*;
import java.lang.ref.*;
import java.util.*;

import org.apache.log4j.*;
import org.dicr.mng.exc.*;
import org.dicr.util.net.*;

/**
 * Класс для учета траффика при помощи iptables.
 * <P>
 * Для подсчета траффика используется отдельная цепочта iptables
 * <code>{@value #DEFAULT_TRAF_CHAIN}</code>, которая должна быть создана в
 * таблице <code>filter</code> и через которую должен проходить весь
 * подсчитываемый траффик. Начальная конфигурация iptables для работы с этим
 * модулем может быть приблизительно следующей:
 * 
 * <PRE>
 * {@value #DEFAULT_IPTABLES_CMD} -N {@value #DEFAULT_TRAF_CHAIN} {@value #DEFAULT_IPTABLES_CMD} -I INPUT -s
 * 192.168.0.0/16 -j {@value #DEFAULT_TRAF_CHAIN} {@value #DEFAULT_IPTABLES_CMD} -I FORWARD -s 192.168.0.0/16 -j
 * {@value #DEFAULT_TRAF_CHAIN} {@value #DEFAULT_IPTABLES_CMD} -I FORWARD -d 192.168.0.0/16 -j
 * {@value #DEFAULT_TRAF_CHAIN} {@value #DEFAULT_IPTABLES_CMD} -I OUTPUT -d 192.168.0.0/16 -j
 * {@value #DEFAULT_TRAF_CHAIN}
 * </PRE>
 * 
 * </P>
 * <P>
 * Все операции по добавлению/удалению/изменению счетчиков выполняются модулем
 * автоматически в цепочке {@value #DEFAULT_TRAF_CHAIN}. Для подсчета траффика
 * на конкретный адрес в цепочку {@value #DEFAULT_TRAF_CHAIN} вставляется 2
 * правила: <CODE>ruleIn</CODE> для входящего и <CODE>ruleOu</CODE> для
 * исходящего.
 * </P>
 * <P>
 * Для блокирования траффика используется цель правила '
 * <CODE>{@value #DEFAULT_TARGET_REJECT}</CODE>', для принятия - '
 * <CODE>{@value #DEFAULT_TARGET_ACCEPT}</CODE>'.
 * </P>
 * <P>
 * Пример подсчета траффика для адреса <code>192.1678.0.1</code>:
 * 
 * <PRE>
 * {@value #DEFAULT_IPTABLES_CMD} -t filter -A {@value #DEFAULT_TRAF_CHAIN} -s 192.168.0.1
 * {@value #DEFAULT_IPTABLES_CMD} -t filter -A {@value #DEFAULT_TRAF_CHAIN} -d 192.168.0.1
 * </PRE>
 * 
 * </P>
 * <P>
 * Пример смены состояния траффика для этого адреса на блокирование:
 * 
 * <PRE>
 * {@value #DEFAULT_IPTABLES_CMD} -t filter -R {@value #DEFAULT_TRAF_CHAIN} 1 -s 192.168.0.1 -j
 * {@value #DEFAULT_TARGET_REJECT} {@value #DEFAULT_IPTABLES_CMD} -t filter -R {@value #DEFAULT_TRAF_CHAIN} 2 -d
 * 192.168.0.1 -j {@value #DEFAULT_TARGET_REJECT}
 * </PRE>
 * 
 * При изменении (добавлении, открытии/закрытии траффика) значение счетчика
 * сбрасывается.
 * </P>
 * <P>
 * Листинг счетчиков выполняется командой:
 * 
 * <PRE>
 * {@value #DEFAULT_IPTABLES_CMD} -nvx -L {@value #DEFAULT_TRAF_CHAIN} --line-numbers
 * </PRE>
 * 
 * </P>
 * <P>
 * Список счетчиков кэшируется в течении <CODE>cacheTimeout</CODE> времени. При
 * листинге счетчиков <CODE>listCounters()</CODE> данные возвращаются из кэша,
 * при условии:
 * 
 * <PRE>
 * System.currentTimeMillis() - getDataTime() &lt; getCachePeriod();
 * </PRE>
 * 
 * При листинге траффика <CODE>listTrafs()</CODE> со сбросом или нет, а также
 * при модификации (добавлении/изменении/удалении) счетчиков кэш сбрасывается.
 * </P>
 * 
 * @author Igor A Tarasov &lt;java@dicr.org&gt;
 * @version 1.0
 * @see IPTablesData
 */
public class IPTables {

	/** Logger */
	private static final Logger log = Logger.getLogger(IPTables.class);

	/** Single instance */
	private static final IPTables iptables = new IPTables();

	/** Default path for iptables binary. */
	public static final String DEFAULT_IPTABLES_CMD = "/sbin/iptables";

	/**
	 * Traffic accounting chain name. This is the special chain in filter table.
	 * Default '<CODE>{@value #DEFAULT_TRAF_CHAIN}</CODE>'.
	 */
	public static final String DEFAULT_TRAF_CHAIN = "traf";

	/** Default accept target name. */
	public static final String DEFAULT_TARGET_ACCEPT = "";

	/** Default reject target */
	public static final String DEFAULT_TARGET_REJECT = "REJECT";

	/** Cache */
	private SoftReference<Map<Net, IPTablesData>> cacheReference = null;

	/** Cache timeout in milliseconds */
	private long cacheTimeout = 5000;

	/** time of data in cache */
	private long cacheTime = 0;

	/**
	 * Constructor
	 */
	protected IPTables() {
		super();
	}

	/**
	 * Return instance
	 * 
	 * @return instance of IPtables
	 */
	public static final IPTables getInstance() {
		return IPTables.iptables;
	}

	/**
	 * Set cache timeout.
	 * 
	 * @param seconds
	 *            cache tgimeout in seconds
	 */
	public final synchronized void setCacheTimeout(final long seconds) {
		if (seconds < 0) throw new IllegalArgumentException("seconds=" + seconds);
		this.cacheTimeout = seconds * 1000;
		IPTables.log.info("configured cache timeout: " + seconds + "s");
	}

	/**
	 * Return cache timeout.
	 * 
	 * @return cache data timeout
	 */
	public final synchronized long getCacheTimeout() {
		return this.cacheTimeout;
	}

	/**
	 * Return cache data time
	 * 
	 * @return time, in which cache data last updated
	 */
	public final synchronized long getCacheTime() {
		return this.cacheTime;
	}

	/**
	 * Translate command to string
	 * 
	 * @param command
	 * @return command string
	 */
	private static final String commandToString(final List<String> command) {
		if (command == null) throw new IllegalArgumentException("null command");
		final StringBuilder sb = new StringBuilder();
		for (final String cp : command) {
			sb.append(cp).append(' ');
		}
		return sb.toString();
	}

	/**
	 * Run iptables command.
	 * 
	 * @param arguments
	 *            iptables arguments. (-t filter will added automatically)
	 * @return command output
	 * @throws CommandException
	 *             error command invocation
	 */
	private final synchronized List<String> runIptables(final List<String> arguments) throws CommandException {
		if (arguments == null || arguments.isEmpty()) throw new IllegalArgumentException("empty command");
		// prepare command string to execute
		final List<String> cmd = new ArrayList<>();
		cmd.add(IPTables.DEFAULT_IPTABLES_CMD);
		cmd.add("-t");
		cmd.add("filter");
		cmd.addAll(arguments);
		IPTables.log.trace("running command: " + IPTables.commandToString(cmd));
		final List<String> out = new ArrayList<>();
		BufferedReader in = null;
		Process p = null;
		try {
			// build process
			final ProcessBuilder pb = new ProcessBuilder(cmd);
			pb.redirectErrorStream(true);
			// start command process
			p = pb.start();
			// read process output
			in = new BufferedReader(new InputStreamReader(p.getInputStream()));
			String line = null;
			while ((line = in.readLine()) != null) {
				out.add(line);
			}
			// check exit vaue
			p.waitFor();
			if (p.exitValue() != 0) throw new CommandException(cmd, out.toString());
		} catch (final IOException ioex) {
			IPTables.log.error("error running iptables", ioex);
			throw new CommandException(cmd, ioex);
		} catch (final InterruptedException ex) {
			IPTables.log.error("iptables execution interrupted", ex);
			throw new CommandException(cmd, ex);
		} finally {
			try {
				if (in != null) in.close();
			} catch (final IOException ex) {
				// NOP
			}
			try {
				if (p != null) p.destroy();
			} catch (final Exception ex) {
				// NOP
			}
		}
		return out;
	}

	/**
	 * List traffic rules
	 * 
	 * @param doReset
	 *            если true, то после листинга счетчики будут обнудены.
	 * @param forceRefresh
	 *            если true, то данные не берутся из кэша.
	 * @throws CommandException
	 *             ошибка выполнения команды iptables.
	 * @return данные счетчиков.
	 */
	private final synchronized Map<Net, IPTablesData> listRules(final boolean doReset, boolean forceRefresh)
			throws CommandException {
		Map<Net, IPTablesData> entries = null;
		// check cache
		if (this.cacheReference != null) entries = this.cacheReference.get();
		if (!forceRefresh && entries != null && System.currentTimeMillis() - this.cacheTime < this.cacheTimeout) return entries;
		// reset cache
		entries = new HashMap<>();
		this.cacheReference = null;
		this.cacheTime = 0;
		// run list command
		final List<String> cmd = new ArrayList<>();
		cmd.add("-nvx");
		cmd.add("--line-numbers");
		if (doReset) cmd.add("-Z");
		cmd.add("-L");
		cmd.add(IPTables.DEFAULT_TRAF_CHAIN);
		final List<String> out = this.runIptables(cmd);
		// parse out
		for (String line : out) {
			try {
				line = line.trim();
				if (line.length() < 80 || line.startsWith("Chain ") || line.startsWith("num ")
						|| line.startsWith("Zeroing ")) continue;
				// parse line
				final int num = Integer.parseInt(line.substring(0, 0 + 4).trim());
				final long bytes = Long.parseLong(line.substring(14, 14 + 8).trim());
				final boolean active = !line.substring(23, 23 + 8).trim()
						.equalsIgnoreCase(IPTables.DEFAULT_TARGET_REJECT);
				final String src = line.substring(58, 58 + 19).trim();
				final String dst = line.substring(79).trim();
				// direction of counter
				final boolean dirsrc = dst.startsWith("0.0.0.0");
				final Net net = new Net(dirsrc ? src : dst);
				IPTablesData counter = entries.get(net);
				if (counter == null) counter = new IPTablesData(net, active);
				// update bytes
				if (dirsrc) {
					counter.setBytesIn(bytes);
					counter.setRuleIn(num);
				} else {
					counter.setBytesOu(bytes);
					counter.setRuleOu(num);
				}
				entries.put(net, counter);
			} catch (final IncorrectAddressException ex) {
				IPTables.log.warn("error parsing address in line: " + line, ex);
			} catch (final NumberFormatException ex) {
				IPTables.log.warn("error parsing line: " + line, ex);
			}
		}
		// update cache
		this.cacheReference = new SoftReference<>(entries);
		this.cacheTime = System.currentTimeMillis();
		IPTables.log.trace("cache updated");
		return entries;
	}

	/**
	 * Set traffic rule. Set rule for accept or reject traffic.
	 * <P>
	 * Не проверяет наличие счетчика, поэтому недоступно для пользователя.
	 * </P>
	 * 
	 * @param ip
	 *            числовое значение IP-адреса.
	 * @param ruleIn
	 *            номер правила во входящей цепочке.
	 * @param ruleOu
	 *            номер правила в исходящей цепочке.
	 * @param opened
	 *            true для разрешения траффика.
	 * @throws CommandException
	 *             ошибка выполнения команды iptables
	 */
	private final synchronized void setTrafficRule(final Net address, final int ruleIn, final int ruleOu,
			final boolean active) throws CommandException {
		if (address == null) throw new IllegalArgumentException("null net");
		final List<String> cmd = new ArrayList<>();
		cmd.add("-R");
		cmd.add(IPTables.DEFAULT_TRAF_CHAIN);
		cmd.add(Integer.toString(ruleIn));
		cmd.add("-s");
		cmd.add(address.toString());
		final String target = active ? IPTables.DEFAULT_TARGET_ACCEPT : IPTables.DEFAULT_TARGET_REJECT;
		if (target.length() > 0) {
			cmd.add("-j");
			cmd.add(target);
		}
		this.runIptables(cmd);
		cmd.set(2, Integer.toString(ruleOu));
		cmd.set(3, "-d");
		this.runIptables(cmd);
		// reset cache
		this.cacheReference = null;
		this.cacheTime = 0;
		IPTables.log.debug("updated traffic rule for '" + address + "' to " + active);
	}

	/**
	 * Add new traffic rule
	 * <P>
	 * Не проверяет наличие дубливанного, поэтому недоступно для пользователя.
	 * </P>
	 * 
	 * @param ip
	 *            числовое значение IP-адреса.
	 * @param opened
	 *            true для разрешения траффика.
	 * @throws CommandException
	 *             ошибка выполнения команды iptables
	 * @see #setTraf
	 * @see #initTrafs
	 */
	private final synchronized void addCounter(final Net address, final boolean active) throws CommandException {
		if (address == null) throw new IllegalArgumentException("null net");
		// run command
		final List<String> cmd = new ArrayList<>();
		cmd.add("-A");
		cmd.add(IPTables.DEFAULT_TRAF_CHAIN);
		cmd.add("-s");
		cmd.add(address.toString());
		final String target = active ? IPTables.DEFAULT_TARGET_ACCEPT : IPTables.DEFAULT_TARGET_REJECT;
		if (target.length() > 0) {
			cmd.add("-j");
			cmd.add(target);
		}
		this.runIptables(cmd);
		cmd.set(2, "-d");
		this.runIptables(cmd);
		// reset cache
		this.cacheReference = null;
		this.cacheTime = 0;
		IPTables.log.debug("added iptables counter for address '" + address + "'");
	}

	/**
	 * Удаляет счетчик.
	 * 
	 * @param ip
	 *            числовое значение IP-адреса.
	 * @param ruleIn
	 *            номер правила во входящей цепочке.
	 * @param ruleOu
	 *            номер правила в исходящей цепочке.
	 * @throws CommandException
	 *             ошибка выполнения команды iptables.
	 */
	private final synchronized void removeCounter(final Net address, final int ruleIn, final int ruleOu)
			throws CommandException {
		if (address == null) throw new IllegalArgumentException("null net");
		final List<String> cmd = new ArrayList<>();
		cmd.add("-D");
		cmd.add(IPTables.DEFAULT_TRAF_CHAIN);
		cmd.add(Integer.toString(ruleIn));
		this.runIptables(cmd);
		cmd.set(2, Integer.toString(ruleOu));
		this.runIptables(cmd);
		// reset cache
		this.cacheTime = 0;
		this.cacheReference = null;
		IPTables.log.debug("removed iptables counter for address '" + address + "'");
	}

	/**
	 * Return counters data.
	 * <P>
	 * Данные возвращаются из кэша, который при необходимости обновляется.
	 * </P>
	 * 
	 * @param reset
	 *            resec counters after listing
	 * @param refresh
	 *            refresh rules (don't use from cache)
	 * @return список и состояние счетчиков.
	 * @throws CommandException
	 *             ошибка выполнения команды iptables.
	 */
	public final Map<Net, IPTablesData> listCounters(final boolean reset, final boolean refresh)
			throws CommandException {
		return Collections.unmodifiableMap(this.listRules(reset, refresh));
	}

	/**
	 * Возвращает состояние счетчиков.
	 * <P>
	 * Кэш принудительно обновляется.
	 * </P>
	 * 
	 * @param reset
	 *            true, если необходимо обнулять счетчики после чтения.
	 * @return данные траффика.
	 * @throws CommandException
	 *             ошибка выполнения команды iptables.
	 */
	public final Map<Net, IPTablesData> getTraf(final boolean reset) throws CommandException {
		return Collections.unmodifiableMap(this.listRules(reset, true));
	}

	/**
	 * Устанавливает счетчик траффика.
	 * <P>
	 * Если счетчика с заданным адресом не существует, то он создается, иначе
	 * его состояние устанавливает в заданное. При смене состояния счетчик байт
	 * сбрасывается в 0.
	 * </P>
	 * 
	 * @param address
	 *            network address
	 * @param active
	 *            true, если траффик разрешен.
	 * @throws CommandException
	 *             ошибка выполнения команды
	 */
	public final synchronized void setCounter(final Net address, final boolean active) throws CommandException {
		if (address == null) throw new IllegalArgumentException("null net");
		final Map<Net, IPTablesData> entries = this.listRules(false, false);
		final IPTablesData counter = entries.get(address);
		if (counter != null) {
			if (counter.isOpened() != active) this.setTrafficRule(address, counter.getRuleIn(), counter.getRuleOu(),
					active);
		} else {
			this.addCounter(address, active);
		}
	}

	/**
	 * Удаляет счетчик.
	 * <P>
	 * Если адрес не существует, то операция не производится и удаление
	 * считается успешным.
	 * </P>
	 * 
	 * @param address
	 *            network address to remove from counters
	 * @throws CommandException
	 *             ошибка выполнения команды
	 */
	public final synchronized void removeCounter(final Net address) throws CommandException {
		if (address == null) throw new IllegalArgumentException("null net");
		final Map<Net, IPTablesData> entries = this.listRules(false, false);
		final IPTablesData counter = entries.get(address);
		if (counter != null) this.removeCounter(address, counter.getRuleIn(), counter.getRuleOu());
	}

	/**
	 * Удаляет все счетчики.
	 * 
	 * @throws CommandException
	 *             ошибка выполнения команды
	 */
	public final synchronized void removeAllCounters() throws CommandException {
		final List<String> cmd = new ArrayList<>();
		cmd.add("-F");
		cmd.add(IPTables.DEFAULT_TRAF_CHAIN);
		this.runIptables(cmd);
		this.cacheTime = 0;
		this.cacheReference = null;
		IPTables.log.info("cleared all iptables counters");
	}

	/**
	 * Инициализирует счетчики iptables.
	 * <P>
	 * Удаляет все существующие счетчики и устанавливает новые.
	 * </P>
	 * 
	 * @param trafs
	 *            список счетчиков для установки.
	 * @throws CommandException
	 *             ошибка выполнения команды.
	 * @see #resetCounters()
	 * @see #setCounter
	 */
	public final synchronized void createCounters(final Set<IPTablesData> trafs) throws CommandException {
		if (trafs == null) throw new IllegalArgumentException("null trafs");
		this.removeAllCounters();
		for (final IPTablesData traf : trafs) {
			this.addCounter(traf.getAddress(), traf.isOpened());
		}
		IPTables.log.info("iptables counters initialized with " + trafs.size() + " addresses");
	}

	/**
	 * Reset traffic counters to zero.
	 * 
	 * @throws CommandException
	 *             command execution error
	 */
	public final synchronized void resetCounters() throws CommandException {
		final List<String> cmd = new ArrayList<>();
		cmd.add("-Z");
		cmd.add(IPTables.DEFAULT_TRAF_CHAIN);
		this.runIptables(cmd);
		this.cacheTime = 0;
		this.cacheReference = null;
		IPTables.log.debug("iptables counters reset");
	}

}
